//
// Created by kenny on 25-11-11.
//
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "mcp_udp.h"

// FFmpeg头文件
#include <ffmpeg/libavformat/avformat.h>
#include <ffmpeg/libavcodec/avcodec.h>
#include <ffmpeg/libavutil/avutil.h>
#include <ffmpeg/libavutil/imgutils.h>

// UDP管理器全局变量
static mcp_udp_t g_udp_manager;
static int g_packet_count = 0;

// UDP接收回调函数（可选，用于调试）
void udp_recv_callback(uint8_t* data, uint16_t length, const char* from_ip, int from_port) {
    printf("Received %d bytes from %s:%d\n", length, from_ip, from_port);
}

// 发送H.264 NALU单元
int send_nalu(const uint8_t* nalu_data, int nalu_size, const char* server_ip, int server_port) {
    if (nalu_size <= 0) return -1;

    // 发送NALU数据
    int ret = mcp_udp_send_to(&g_udp_manager, nalu_data, nalu_size, server_ip, server_port);
    if (ret == 0) {
        g_packet_count++;
        printf("Sent NALU #%d, size: %d bytes\n", g_packet_count, nalu_size);
    }

    return ret;
}

// 查找起始码 (0x00000001 or 0x000001)
int find_start_code(const uint8_t* data, int size, int* start_code_len) {
    if (size < 4) return -1;

    for (int i = 0; i < size - 3; i++) {
        if (data[i] == 0x00 && data[i+1] == 0x00 && data[i+2] == 0x00 && data[i+3] == 0x01) {
            *start_code_len = 4;
            return i;
        }
        if (data[i] == 0x00 && data[i+1] == 0x00 && data[i+2] == 0x01) {
            *start_code_len = 3;
            return i;
        }
    }
    return -1;
}

// 从AVPacket中提取NALU并发送
int extract_and_send_nalus(AVPacket* pkt, const char* server_ip, int server_port) {
    if (!pkt || pkt->size <= 0) return -1;

    uint8_t* data = pkt->data;
    int size = pkt->size;
    int sent_count = 0;

    // 查找第一个起始码
    int start_code_len;
    int pos = find_start_code(data, size, &start_code_len);

    if (pos < 0) {
        // 没有找到起始码，直接发送整个包
        printf("No start code found, sending entire packet (%d bytes)\n", size);
        return send_nalu(data, size, server_ip, server_port);
    }

    // 发送每个NALU
    while (pos >= 0 && size > 0) {
        data += pos + start_code_len;
        size -= pos + start_code_len;

        if (size <= 0) break;

        // 查找下一个起始码
        int next_pos = find_start_code(data, size, &start_code_len);
        int nalu_size;

        if (next_pos >= 0) {
            nalu_size = next_pos;
        } else {
            nalu_size = size;
        }

        if (nalu_size > 0) {
            // 发送NALU
            if (send_nalu(data, nalu_size, server_ip, server_port) == 0) {
                sent_count++;
            }
        }

        pos = next_pos;
    }

    return sent_count;
}

// 提取并发送H.264流
int extract_h264_stream(const char* input_file, const char* server_ip, int server_port) {
    AVFormatContext* fmt_ctx = NULL;
    AVPacket pkt;
    int video_stream_index = -1;
    int ret = 0;

    // 注册FFmpeg所有组件
    av_register_all();

    // 打开输入文件
    if (avformat_open_input(&fmt_ctx, input_file, NULL, NULL) < 0) {
        fprintf(stderr, "Could not open input file: %s\n", input_file);
        return -1;
    }

    // 获取流信息
    if (avformat_find_stream_info(fmt_ctx, NULL) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    // 查找视频流（H.264）
    for (int i = 0; i < fmt_ctx->nb_streams; i++) {
        if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_index = i;
            break;
        }
    }

    if (video_stream_index == -1) {
        fprintf(stderr, "Could not find video stream\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    // 打印视频流信息
    AVCodecParameters* codec_par = fmt_ctx->streams[video_stream_index]->codecpar;
    printf("Video codec: %d (H.264=%d)\n", codec_par->codec_id, AV_CODEC_ID_H264);
    printf("Resolution: %dx%d\n", codec_par->width, codec_par->height);
    printf("Bitrate: %lld\n", codec_par->bit_rate);

    // 初始化UDP
    if (mcp_udp_init(&g_udp_manager, server_ip, server_port, udp_recv_callback) != 0) {
        fprintf(stderr, "UDP initialization failed\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    if (mcp_udp_connect(&g_udp_manager) != 0) {
        fprintf(stderr, "UDP connection failed\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    printf("Start sending H.264 stream to %s:%d\n", server_ip, server_port);

    // 读取并发送数据包
    av_init_packet(&pkt);
    pkt.data = NULL;
    pkt.size = 0;

    while (av_read_frame(fmt_ctx, &pkt) >= 0) {
        if (pkt.stream_index == video_stream_index) {
            // 发送视频数据包
            extract_and_send_nalus(&pkt, server_ip, server_port);

            // 添加延迟以控制发送速率（根据实际需要调整）
            usleep(33000); // 约30fps
        }
        av_packet_unref(&pkt);
    }

    printf("Finished sending H.264 stream. Total packets: %d\n", g_packet_count);

    // 清理资源
    avformat_close_input(&fmt_ctx);
    mcp_udp_close(&g_udp_manager);
    mcp_udp_cleanup(&g_udp_manager);

    return 0;
}

// 简化版本：直接发送AVPacket数据
int send_h264_simple(const char* input_file, const char* server_ip, int server_port) {
    AVFormatContext* fmt_ctx = NULL;
    AVPacket pkt;
    int video_stream_index = -1;
    int ret = 0;

    av_register_all();

    if (avformat_open_input(&fmt_ctx, input_file, NULL, NULL) < 0) {
        fprintf(stderr, "Could not open input file: %s\n", input_file);
        return -1;
    }

    if (avformat_find_stream_info(fmt_ctx, NULL) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    for (int i = 0; i < fmt_ctx->nb_streams; i++) {
        if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_stream_index = i;
            break;
        }
    }

    if (video_stream_index == -1) {
        fprintf(stderr, "Could not find video stream\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    // 初始化UDP
    if (mcp_udp_init(&g_udp_manager, server_ip, server_port, udp_recv_callback) != 0) {
        fprintf(stderr, "UDP initialization failed\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    if (mcp_udp_connect(&g_udp_manager) != 0) {
        fprintf(stderr, "UDP connection failed\n");
        avformat_close_input(&fmt_ctx);
        return -1;
    }

    printf("Start sending H.264 stream to %s:%d\n", server_ip, server_port);

    av_init_packet(&pkt);
    pkt.data = NULL;
    pkt.size = 0;

    while (av_read_frame(fmt_ctx, &pkt) >= 0) {
        if (pkt.stream_index == video_stream_index) {
            // 直接发送整个AVPacket数据
            if (mcp_udp_send_to(&g_udp_manager, pkt.data, pkt.size, server_ip, server_port) == 0) {
                g_packet_count++;
                printf("Sent packet #%d, size: %d bytes\n", g_packet_count, pkt.size);
            }

            // 控制发送速率
            usleep(33000); // 约30fps
        }
        av_packet_unref(&pkt);
    }

    printf("Finished sending H.264 stream. Total packets: %d\n", g_packet_count);

    avformat_close_input(&fmt_ctx);
    mcp_udp_close(&g_udp_manager);
    mcp_udp_cleanup(&g_udp_manager);

    return 0;
}

int main(int argc, char* argv[]) {
    if (argc != 4) {
        printf("Usage: %s <input.mp4> <server_ip> <server_port>\n", argv[0]);
        printf("Example: %s test.mp4 192.168.1.100 12345\n", argv[0]);
        return -1;
    }

    const char* input_file = argv[1];
    const char* server_ip = argv[2];
    int server_port = atoi(argv[3]);

    printf("Extracting H.264 from: %s\n", input_file);
    printf("Sending to: %s:%d\n", server_ip, server_port);

    // 使用简化版本（推荐）
    return send_h264_simple(input_file, server_ip, server_port);

    // 或者使用NALU提取版本
    // return extract_h264_stream(input_file, server_ip, server_port);
}